#include "model_landing_3dof.h"
/**
 *      三自由度模型,弹体速度系和弹体系重合，攻角，侧滑作为等效发动机摆角，方向遵循此方向。见航飞动
 *      x y z vx vy vz theta theta_dot phi dphi_dot
 *      m dv/dt = P*cos(alpha)*cos(beta) - X - mgsin(theta)
 *      m v dtheta/dt = P*sin(alpha) - mgcos(theta)
 *      m v cos(theta) dphi/dt = P*cos(alpha)sin(beta) 
 *      dx/dt = V*cos(theta)cos(phi)
 *      dy/dt = V*sin(theta)
 *      dz/dt = -V*cos(theta)sin(phi)
 *      //dtheta_dot/dt = -P*sin(alpha)*rTB/m*(rG*rG)
 *      //dphi_dot/dt = -P*cos(alpha)sin(beta)*rTB/m*(rG*rG)
 *       
 * 
 * 
 * 
 */

/**
 * @brief 按照给定的初始状态初始化状态和控制量，采用的方法是线性插值，终点到初始点的一条线段
 * 
 * @param X 状态量
 * @param U 输入量
 */
void model_landing_3dof::initialize(Eigen::Matrix<double, n_states, K> &X, Eigen::Matrix<double, n_inputs, K> &U) {

    X.setZero();
    U.setZero();
    for (size_t k = 0; k < K; k++) {
        double alpha2 = double(k)/K;
        //double alpha1 = 1. - alpha2;

        // X(0,k) = rx_init * alpha1;
        // X(1,k) = ry_init * alpha1;
        // X(2,k) = vx_init * alpha1;
        // X(3,k) = vy_init * alpha1;
        // X(4,k) = theta_init * alpha1;
        // X(5,k) = dtheta_init * alpha1;
        X(0,k) = (rx_final - rx_init) * alpha2 + rx_init;
        X(1,k) = (ry_final - ry_init) * alpha2 + ry_init;
        X(2,k) = (rz_final - rz_init) * alpha2 + rz_init;

        X(3,k) = (vx_final - vx_init) * alpha2 + vx_init;
        X(4,k) = (vy_final - vy_init) * alpha2 + vy_init;
        X(5,k) = (vz_final - vz_init) * alpha2 + vz_init;

        X(6,k) = (theta_final - theta_init) * alpha2 + theta_init;
        X(7,k) = (dtheta_final - dtheta_init) * alpha2 + dtheta_init;
        X(8,k) = (phi_final - phi_init) * alpha2 + phi_init;
        X(9,k) = (dphi_final - dphi_init) * alpha2 + dphi_init;
        

        U(0,k) = 1./TWR;
        U(1,k) = 0.0;
        U(2,k) = 0.0;
    }

}

// model_landing_3dof::StateVector model_landing_3dof::ode(const StateVector &x, const ControlVector &u) {
//     /*
//         考虑在三自由度中，竖直方向为y轴，水平方向为x轴，整个角度(包括推力方向与火箭夹角
//         火箭纵轴与竖直方向夹角)为顺时针为正，动力学力矩方向也是顺时针为正
//     */
//     const double throttle = u(0, 0);
//     const double gimbalAngle = u(1, 0);
//     const double vx = x(2, 0);
//     const double vy = x(3, 0);
//     const double theta = x(4, 0);
//     const double dtheta = x(5, 0);
//     const double x0 = TWR*g*throttle;
//     const double x1 = gimbalAngle + theta;
//     //增加气动力项
//     double f_d =  C_D*(vx*vx+vy*vy);
//     //double f_d_theta = PI/2 - atan2(-vy,-vx);
//     double f_d_theta = atan2(-vx,-vy);
//     if(f_d_theta > PI/2) f_d_theta = f_d_theta - PI;
//     else if(f_d_theta < -PI/2) f_d_theta = f_d_theta + PI;

//     StateVector f;
//     f(0, 0) = vx;
//     f(1, 0) = vy;
//     f(2, 0) = x0*sin(x1) + f_d*sin(f_d_theta);
//     f(3, 0) = g*(TWR*throttle*cos(x1) - 1) + f_d*cos(f_d_theta);
//     f(4, 0) = dtheta;
//     f(5, 0) = -rTB*x0*sin(gimbalAngle)/(rG*rG);


//     return f;
// }

model_landing_3dof::StateVector model_landing_3dof::ode(const StateVector &x, const ControlVector &u)
{
    const double throttle = u(0, 0);
    const double alpha = u(1, 0);
    const double beta = u(2, 0);
    const double vx = x(3, 0);
    const double vy = x(4, 0);
    const double vz = x(5, 0);
    const double V = sqrt(vx*vx + vy*vy + vz*vz);
    const double theta = x(6, 0);
    //const double dtheta = x(7, 0);
    const double phi = x(8, 0);
    //const double dphi = x(9, 0);
    const double X_m = C_D*(vx*vx + vy*vy + vz*vz);
    const double P = TWR*g*throttle;

    double dv_dt = P*cos(alpha)*cos(beta) - X_m - g*sin(theta);
    double dtheta_dt = P*sin(alpha)/V - g*cos(theta)/V;
    double dphi_dt = P*cos(alpha)*sin(beta)/(V*cos(theta)); 
    double dx_dt = V*cos(theta)*cos(phi);
    double dy_dt = V*sin(theta);
    double dz_dt = -V*cos(theta)*sin(phi);
    double ddx_dt = dv_dt*cos(theta)*cos(phi);
    double ddy_dt = dv_dt*sin(theta);
    double ddz_dt = -dv_dt*cos(theta)*sin(phi);
    double dtheta_dot_dt = -P*sin(alpha)*rTB/(rG*rG);
    double dphi_dot_dt = -P*cos(alpha)*sin(beta)*rTB/(rG*rG);

    StateVector f;
    f(0,0) = dx_dt;
    f(1,0) = dy_dt;
    f(2,0) = dz_dt;
    f(3,0) = ddx_dt;
    f(4,0) = ddy_dt;
    f(5,0) = ddz_dt;
    f(6,0) = dtheta_dt;
    f(7,0) = dtheta_dot_dt;
    f(8,0) = dphi_dt;
    f(9,0) = dphi_dot_dt;

    return f;
}
// model_landing_3dof::StateMatrix model_landing_3dof::state_jacobian(const StateVector &x, const ControlVector &u) {

//     const double throttle = u(0, 0);
//     const double gimbalAngle = u(1, 0);
//     const double theta = x(4, 0);
//     const double x0 = TWR*g*throttle;
//     const double x1 = gimbalAngle + theta;
//     //增加气动力项
//     const double vx = x(2, 0);
//     const double vy = x(3, 0);
//     //double f_d_theta = PI/2 - atan2(-vy,-vx);
//     double f_d_theta = atan2(-vx,-vy);
//     if(f_d_theta > PI/2) f_d_theta = f_d_theta - PI;
//     else if(f_d_theta < -PI/2) f_d_theta = f_d_theta + PI;

//     StateMatrix A;
//     A.setZero();
//     A(0, 2) = 1;
//     A(1, 3) = 1;
//     A(2, 2) = 2*C_D*vx*sin(f_d_theta) - C_D*vy*cos(f_d_theta); //v_x_dot 对 v_x
//     A(2, 3) = 2*C_D*vy*sin(f_d_theta) + C_D*vx*cos(f_d_theta);//v_x_dot 对 v_y
//     // A(2, 2) = 2*C_D*vx*sin(f_d_theta) + C_D*vy*cos(f_d_theta); //v_x_dot 对 v_x
//     // A(2, 3) = 2*C_D*vy*sin(f_d_theta) - C_D*vx*cos(f_d_theta);//v_x_dot 对 v_y
//     A(2, 4) = x0*cos(x1);
//     A(3, 2) =  + C_D*vy*sin(f_d_theta) + 2*C_D*vx*cos(f_d_theta);//v_y_dot 对 v_x
//     A(3, 3) =  - C_D*vx*sin(f_d_theta) + 2*C_D*vy*cos(f_d_theta);//v_y_dot 对 v_x
//     // A(3, 2) =  - C_D*vy*sin(f_d_theta) + 2*C_D*vx*cos(f_d_theta);//v_y_dot 对 v_x
//     // A(3, 3) =  C_D*vx*sin(f_d_theta) + 2*C_D*vy*cos(f_d_theta);//v_y_dot 对 v_x
//     A(3, 4) = -x0*sin(x1);
//     A(4, 5) = 1;


//     return A;
// }
model_landing_3dof::StateMatrix model_landing_3dof::state_jacobian(const StateVector &x, const ControlVector &u)
{
    //x y z vx vy vz theta theta_dot phi dphi_dot
    // double x = x(0, 0);
    // double y = x(1, 0);
    // double z = x(2, 0);
    double vx = x(3, 0);
    double vy = x(4, 0);
    double vz = x(5, 0);
    double theta = x(6, 0);
    //double theta_dot = x(7, 0);
    double phi = x(8, 0);
    //double phi_dot = x(9, 0);
    double throttle = u(0, 0);
    double alpha = u(1, 0);
    double beta = u(2, 0);
    double X_m = C_D*(vx*vx + vy*vy + vz*vz);
    double P = TWR*g*throttle;
    double v = sqrt(vx*vx + vy*vy + vz*vz);

    model_landing_3dof::StateMatrix jacobian;
    jacobian.setZero();
    jacobian(0, 3) = 1;
    jacobian(1, 4) = 1;
    jacobian(2, 5) = 1;
    jacobian(3, 6) = cos(phi) * sin(theta) * (X_m + g * sin(theta) - P * cos(alpha) * cos(beta)) - g * cos(phi) * cos(theta) * cos(theta);
    jacobian(3, 8) = cos(theta) * sin(phi) * (X_m + g * sin(theta) - P * cos(alpha) * cos(beta));
    jacobian(4, 6) = -cos(theta) * (X_m + g * sin(theta) - P * cos(alpha) * cos(beta)) - g * cos(theta) * sin(theta);
    jacobian(5, 6) = g * cos(theta) * cos(theta) * sin(phi) - sin(phi) * sin(theta) * (X_m + g * sin(theta) - P * cos(alpha) * cos(beta));
    jacobian(5, 8) = cos(phi) * cos(theta) * (X_m + g * sin(theta) - P * cos(alpha) * cos(beta));
    jacobian(6, 3) = (g * vx * cos(theta)) / (v*v*v) - (P * vx * sin(alpha)) / (v*v*v);
    jacobian(6, 4) = (g * vy * cos(theta)) / (v*v*v) - (P * vy * sin(alpha)) / (v*v*v);
    jacobian(6, 5) = (g * vz * cos(theta)) / (v*v*v) - (P * vz * sin(alpha)) / (v*v*v);
    jacobian(6, 6) = (g * sin(theta)) / v;
    jacobian(8, 3) = -(P * vx * cos(alpha) * sin(beta)) / (cos(theta) * (v*v*v));
    jacobian(8, 4) = -(P * vy * cos(alpha) * sin(beta)) / (cos(theta) * (v*v*v));
    jacobian(8, 5) = -(P * vz * cos(alpha) * sin(beta)) / (cos(theta) * (v*v*v));
    jacobian(8, 6) = (P * cos(alpha) * sin(beta) * sin(theta)) / (cos(theta) * cos(theta) * v);

    return jacobian;
}

model_landing_3dof::ControlMatrix model_landing_3dof::control_jacobian(const StateVector &x, const ControlVector &u)
{
    double vx = x(3, 0);
    double vy = x(4, 0);
    double vz = x(5, 0);
    double theta = x(6, 0);
    double phi = x(8, 0);
    double throttle = u(0, 0);
    double alpha = u(1, 0);
    double beta = u(2, 0);
    //double X_m = C_D*(vx*vx + vy*vy + vz*vz);
    double P = TWR*g*throttle;
    double v = sqrt(vx*vx + vy*vy + vz*vz);

    model_landing_3dof::ControlMatrix jacobian;
    jacobian.setZero();
    jacobian(3, 0) = cos(alpha) * cos(beta) * cos(phi) * cos(theta);
    jacobian(3, 1) = -P * cos(beta) * cos(phi) * sin(alpha) * cos(theta);
    jacobian(3, 2) = -P * cos(alpha) * cos(phi) * sin(beta) * cos(theta);
    jacobian(4, 0) = cos(alpha) * cos(beta) * sin(theta);
    jacobian(4, 1) = -P * cos(beta) * sin(alpha) * sin(theta);
    jacobian(4, 2) = -P * cos(alpha) * sin(beta) * sin(theta);
    jacobian(5, 0) = -cos(alpha) * cos(beta) * cos(theta) * sin(phi);
    jacobian(5, 1) = P * cos(beta) * sin(alpha) * cos(theta) * sin(phi);
    jacobian(5, 2) = P * cos(alpha) * sin(beta) * cos(theta) * sin(phi);
    jacobian(6, 0) = sin(alpha) / v;
    jacobian(6, 1) = (P * cos(alpha)) / v;
    jacobian(6, 2) = 0;
    jacobian(7, 0) = -(rTB * sin(alpha)) / (rG * rG);
    jacobian(7, 1) = -(P * rTB * cos(alpha)) / (rG * rG);
    jacobian(7, 2) = 0;
    jacobian(8, 0) = (cos(alpha) * sin(beta)) / (cos(theta) * v);
    jacobian(8, 1) = -(P * sin(alpha) * sin(beta)) / (cos(theta) * v);
    jacobian(8, 2) = (P * cos(alpha) * cos(beta)) / (cos(theta) * v);
    jacobian(9, 0) = -(rTB * cos(alpha) * sin(beta)) / (rG * rG);
    jacobian(9, 1) = (P * rTB * sin(alpha) * sin(beta)) / (rG * rG);
    jacobian(9, 2) = -(P * rTB * cos(alpha) * cos(beta)) / (rG * rG);

    return jacobian;
}
// model_landing_3dof::ControlMatrix model_landing_3dof::control_jacobian(const StateVector &x, const ControlVector &u) {

//     const double throttle = u(0, 0);
//     const double gimbalAngle = u(1, 0);
//     const double theta = x(4, 0);
//     const double x0 = TWR*g;
//     const double x1 = gimbalAngle + theta;
//     const double x2 = x0*sin(x1);
//     const double x3 = x0*cos(x1);
//     const double x4 = TWR*g*rTB/(rG*rG);

//     ControlMatrix B;
//     B.setZero();
//     B(2, 0) = x2;
//     B(2, 1) = throttle*x3;
//     B(3, 0) = x3;
//     B(3, 1) = -throttle*x2;
//     B(5, 0) = -x4*sin(gimbalAngle);
//     B(5, 1) = -throttle*x4*cos(gimbalAngle);


//     return B; 
// }

void model_landing_3dof::add_application_constraints(
    optimization_problem::SecondOrderConeProgram &socp,
    Eigen::Matrix<double, n_states, K> &X0,
    Eigen::Matrix<double, n_inputs, K> &U0
) {

    auto var = [&](const string &name, const vector<size_t> &indices){ return socp.get_variable(name,indices); };
    auto param = [](double &param_value){ return optimization_problem::Parameter(&param_value); };
    
    // initial state
    socp.add_constraint( (-1.0) * var("X", {0, 0}) + param(rx_init) == 0.0 );
    socp.add_constraint( (-1.0) * var("X", {1, 0}) + param(ry_init) == 0.0 );
    socp.add_constraint( (-1.0) * var("X", {2, 0}) + param(rz_init) == 0.0 );
    socp.add_constraint( (-1.0) * var("X", {3, 0}) + param(vx_init) == 0.0 );
    socp.add_constraint( (-1.0) * var("X", {4, 0}) + param(vy_init) == 0.0 );
    socp.add_constraint( (-1.0) * var("X", {5, 0}) + param(vz_init) == 0.0 );
    socp.add_constraint( (-1.0) * var("X", {6, 0}) + param(theta_init) == 0.0 );
    socp.add_constraint( (-1.0) * var("X", {7, 0}) + param(dtheta_init) == 0.0 );
    socp.add_constraint( (-1.0) * var("X", {8, 0}) + param(phi_init) == 0.0 );
    socp.add_constraint( (-1.0) * var("X", {9, 0}) + param(dphi_init) == 0.0 );


    // final state
    socp.add_constraint( (1.0) * var("X", {0, K-1}) == 0.0 );
    socp.add_constraint( (1.0) * var("X", {1, K-1}) == 0.0 );
    socp.add_constraint( (1.0) * var("X", {2, K-1}) == 0.0 );
    socp.add_constraint( (1.0) * var("X", {3, K-1}) == 0.0 );
    socp.add_constraint( (1.0) * var("X", {4, K-1}) == 0.0 );
    socp.add_constraint( (1.0) * var("X", {5, K-1}) == 0.0 );
    socp.add_constraint( (-1.0) * var("X", {6, K-1}) + PI/2 == 0.0);
    socp.add_constraint( (1.0) * var("X", {7, K-1}) == 0.0 );
    //socp.add_constraint( (1.0) * var("X", {8, K-1}) == 0.0 );
    socp.add_constraint( (1.0) * var("X", {9, K-1}) == 0.0 );

    // glide slope cone 下滑道约束
    const double tan_gamma_gs = tan(max_glide_slope_angle);
    for (size_t k = 0; k < K; ++k) {
        socp.add_constraint(
            optimization_problem::norm2({ 
                (1.0) * var("X", {2, k}),  // z
                (1.0) * var("X", {3, k})   // vz
            })
            <= (1.0 / tan_gamma_gs) * var("X", {1, k})  // x
        );
    }
    // const double tan_gamma_gs = tan(max_glide_slope_angle);
    // for (size_t k = 0; k < K-1; ++k) {
    //     socp.add_constraint(
    //         optimization_problem::norm2({ 
    //             (1.0) * var("X", {0, k})
    //         })
    //         <= 1/tan_gamma_gs * var("X", {1, k}) 
    //     );
    //     //socp.add_constraint(1/tan_gamma_gs * var("X", {1, k})  + (-1.0) * var("X", {0, k}) >= 0.0);
    // }
    //地面上飞行约束
    for (size_t k = 0; k < K; ++k) {
        socp.add_constraint((1.0) * var("X", {1, k}) >= 0.0);
    }
    // TODO
    //姿态约束
    // for (size_t k = 0; k < K; ++k) {
    //     socp.add_constraint((-1.0) * var("X", {4, k}) + max_attitude_angle >= 0.0);
    //     socp.add_constraint((1.0) * var("X", {4, k}) + max_attitude_angle >= 0.0);
    // }

    for (size_t k = 0; k < K; ++k) {

        // throttle control constraints
        socp.add_constraint( ( 1.0) * var("U", {0, k}) + (-0.1) >= (0.0) );
        socp.add_constraint( (-1.0) * var("U", {0, k}) + (1.0) >= (0.0) );

        // gimbal control constraints
        if(k < 0.8 * K) {
            socp.add_constraint( ( 1.0) * var("U", {1, k}) + (max_gimbal_angle) >= (0.0) );
            socp.add_constraint( (-1.0) * var("U", {1, k}) + (max_gimbal_angle) >= (0.0) );
        } else {
            socp.add_constraint( ( 1.0) * var("U", {1, k}) + (0.1 * max_gimbal_angle) >= (0.0) );
            socp.add_constraint( (-1.0) * var("U", {1, k}) + (0.1 * max_gimbal_angle) >= (0.0) );
        }
        if(k < 0.8 * K) {
            socp.add_constraint( ( 1.0) * var("U", {2, k}) + (max_gimbal_angle) >= (0.0) );
            socp.add_constraint( (-1.0) * var("U", {2, k}) + (max_gimbal_angle) >= (0.0) );
        } else {
            socp.add_constraint( ( 1.0) * var("U", {2, k}) + (0.1 * max_gimbal_angle) >= (0.0) );
            socp.add_constraint( (-1.0) * var("U", {2, k}) + (0.1 * max_gimbal_angle) >= (0.0) );
        }
    }
}

// void model_landing_3dof::add_application_constraints(
//     optimization_problem::SecondOrderConeProgram &socp,
//     Eigen::Matrix<double, n_states, K> &X0,
//     Eigen::Matrix<double, n_inputs, K> &U0
// ) {

//     auto var = [&](const string &name, const vector<size_t> &indices){ return socp.get_variable(name,indices); };
//     auto param = [](double &param_value){ return optimization_problem::Parameter(&param_value); };
    
//     // initial state
//     socp.add_constraint( (-1.0) * var("X", {0, 0}) + param(rx_init) == 0.0 );
//     socp.add_constraint( (-1.0) * var("X", {1, 0}) + param(ry_init) == 0.0 );
//     socp.add_constraint( (-1.0) * var("X", {2, 0}) + param(vx_init) == 0.0 );
//     socp.add_constraint( (-1.0) * var("X", {3, 0}) + param(vy_init) == 0.0 );
//     socp.add_constraint( (-1.0) * var("X", {4, 0}) + param(theta_init) == 0.0 );
//     socp.add_constraint( (-1.0) * var("X", {5, 0}) + param(dtheta_init) == 0.0 );


//     // final state
//     socp.add_constraint( (1.0) * var("X", {0, K-1}) == 0.0 );
//     socp.add_constraint( (1.0) * var("X", {1, K-1}) == 0.0 );
//     socp.add_constraint( (1.0) * var("X", {2, K-1}) == 0.0 );
//     socp.add_constraint( (1.0) * var("X", {3, K-1}) == 0.0 );
//     socp.add_constraint( (1.0) * var("X", {4, K-1}) == 0.0 );
//     socp.add_constraint( (1.0) * var("X", {5, K-1}) == 0.0 );

//     // glide slope cone
//     // const double tan_gamma_gs = tan(max_glide_slope_angle);
//     // for (size_t k = 0; k < K; ++k) {
//     //     socp.add_constraint(
//     //         optimization_problem::norm2({ 
//     //             (1.0) * var("X", {2, k}),  // z
//     //             (1.0) * var("X", {3, k})   // vz
//     //         })
//     //         <= (1.0 / tan_gamma_gs) * var("X", {1, k})  // x
//     //     );
//     // }
//     const double tan_gamma_gs = tan(max_glide_slope_angle);
//     for (size_t k = 0; k < K-1; ++k) {
//         socp.add_constraint(
//             optimization_problem::norm2({ 
//                 (1.0) * var("X", {0, k})
//             })
//             <= 1/tan_gamma_gs * var("X", {1, k}) 
//         );
//         //socp.add_constraint(1/tan_gamma_gs * var("X", {1, k})  + (-1.0) * var("X", {0, k}) >= 0.0);
//     }
//     //地面上飞行约束
//     for (size_t k = 0; k < K; ++k) {
//         socp.add_constraint((1.0) * var("X", {1, k}) >= 0.0);
//     }
//     // TODO
//     //姿态约束
//     // for (size_t k = 0; k < K; ++k) {
//     //     socp.add_constraint((-1.0) * var("X", {4, k}) + max_attitude_angle >= 0.0);
//     //     socp.add_constraint((1.0) * var("X", {4, k}) + max_attitude_angle >= 0.0);
//     // }

//     for (size_t k = 0; k < K; ++k) {

//         // throttle control constraints
//         socp.add_constraint( ( 1.0) * var("U", {0, k}) + (-0.1) >= (0.0) );
//         socp.add_constraint( (-1.0) * var("U", {0, k}) + (1.0) >= (0.0) );

//         // gimbal control constraints
//         if(k < 0.8 * K) {
//             socp.add_constraint( ( 1.0) * var("U", {1, k}) + (max_gimbal_angle) >= (0.0) );
//             socp.add_constraint( (-1.0) * var("U", {1, k}) + (max_gimbal_angle) >= (0.0) );
//         } else {
//             socp.add_constraint( ( 1.0) * var("U", {1, k}) + (0.1 * max_gimbal_angle) >= (0.0) );
//             socp.add_constraint( (-1.0) * var("U", {1, k}) + (0.1 * max_gimbal_angle) >= (0.0) );
//         }
//     }
// }

model_landing_3dof::StateVector model_landing_3dof::get_random_state() {
    StateVector X;
    X.setRandom();
    X *= 10;
    return X;
}

model_landing_3dof::ControlVector model_landing_3dof::get_random_input() {
    ControlVector U;
    U.setRandom();
    return U;
}